import numpy as np
import numexpr as ne # to speed up the computations
from PIL import Image, ImageEnhance
from torchvision.datasets import ImageFolder
import torchvision.transforms as transforms


def get_folders():
    
    data_dir = './data/'
    
    train_transform = transforms.Compose([
        transforms.Resize(384),
        transforms.RandomCrop(256),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        ),
    ])
    
    val_transform = transforms.Compose([
        transforms.Resize(384),
        transforms.RandomCrop(256),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        ),
    ])
    
    train_folder = ImageFolder(data_dir + 'train', train_transform)
    val_folder = ImageFolder(data_dir + 'val', val_transform)
    return train_folder, val_folder


# folder name to index: class_to_idx
def get_class_weights(class_to_idx):
    
    # folder name to class name
    decode = np.load('../train_val_split/decode.npy')[()]
    # in the other direction
    encode = {decode[k]: k for k in decode}
    # number of samples in each class
    class_counts = np.load('../preprocessing_utils/class_counts.npy')[()]

    class_counts = {encode[k]: class_counts[k] for k in class_counts}
    # class index to the number of samples in this class
    class_counts = {class_to_idx[str(k)]: class_counts[k] for k in class_counts}

    w = np.zeros((256,))
    for k in class_counts:
        w[k] = class_counts[k]

    w = 1.0/w
    return w, decode



